#! /usr/bin/env python
# -*- coding: utf-8 -*-
import json
import time
import os
import os.path
import sys
import zipfile
from datetime import date


class PepsDownload:
    def __init__(self, location=None,
                 auth=None,
                 write_dir='.',
                 collection='S2',
                 product_type="",
                 sensor_mode="",
                 no_download=False,
                 start_date=None,
                 tile=None,
                 lat=None,
                 lon=None,
                 latmin=None,
                 latmax=None,
                 lonmin=None,
                 lonmax=None,
                 orbit=None,
                 end_date='9999-01-01',
                 search_json_file=None,
                #  windows=None,
                 windows=True,
                 clouds=100,
                 sat=None,
                 extract=None):
        """
        利用PEPS下载哨兵数据
        :param location: town name (pick one which is not too frequent to avoid confusions)
        :param auth: Peps account and password file
        :param write_dir: Path where the products should be downloaded
        :param collection: Collection within theia collections", choices=['S1', 'S2', 'S2ST', 'S3']
        :param product_type: GRD, SLC, OCN (for S1) | S2MSI1C S2MSI2A S2MSI2Ap (for S2)
        :param sensor_mode: EW, IW , SM, WV (for S1) | INS-NOBS, INS-RAW (for S3)
        :param no_download: Do not download products, just print curl command
        :param start_date: start date, fmt('2015-12-22')
        :param tile: Sentinel-2 tile number
        :param lat: latitude in decimal degrees
        :param lon: longitude in decimal degrees
        :param latmin: min latitude in decimal degrees
        :param latmax: max latitude in decimal degrees
        :param lonmin: min longitude in decimal degrees
        :param lonmax: max longitude in decimal degrees
        :param orbit: Orbit Path number
        :param end_date: end date, fmt('2015-12-23')
        :param search_json_file: Output search JSON filename
        :param windows: For windows usage
        :param clouds: Maximum cloud coverage
        :param sat: S1A,S1B,S2A,S2B,S3A,S3B
        :param extract: Extract and remove zip file after download
        """
        # 设置下载参数
        self.location = location
        self.auth = auth
        self.write_dir = write_dir
        self.collection = collection
        self.product_type = product_type
        self.sensor_mode = sensor_mode
        self.no_download = no_download
        self.start_date = start_date
        self.tile = tile
        self.lat = lat
        self.lon = lon
        self.latmin = latmin
        self.latmax = latmax
        self.lonmin = lonmin
        self.lonmax = lonmax
        self.orbit = orbit
        self.end_date = end_date
        self.search_json_file = search_json_file
        self.windows = windows
        self.clouds = clouds
        self.sat = sat
        self.extract = extract
        # 下载文件
        self.downloader()

    ###########################################################################
    def check_rename(self, tmpfile, prodsize, prod):
        print(os.path.getsize(tmpfile), prodsize)
        if os.path.getsize(tmpfile) != prodsize:
            with open(tmpfile) as f_tmp:
                try:
                    tmp_data = json.load(f_tmp)
                    print("Result is a json file (might come from a wrong password file)")
                    print(tmp_data)
                    sys.exit(-1)
                except ValueError:
                    print("\ndownload was not complete, tmp file removed")
                    os.remove(tmpfile)
                    return

        zfile = "%s/%s.zip" % (self.write_dir, prod)
        os.system('mv -f %s %s ' % (tmpfile, zfile))

        # Unzip file
        if self.extract and os.path.exists(zfile):
            try:
                with zipfile.ZipFile(zfile, 'r') as zf:
                    safename = zf.namelist()[0].replace('/', '')
                    zf.extractall(self.write_dir)
                safedir = os.path.join(self.write_dir, safename)
                if not os.path.isdir(safedir):
                    raise Exception('Unzipped directory not found: ', zfile)

            except Exception as e:
                print(e)
                print('Could not unzip file: ' + zfile)
                os.remove(zfile)
                print('Zip file removed.')
                return

            else:
                print('product saved as : ' + safedir)
                os.remove(zfile)
                return

        print("product saved as : " + zfile)

    def parse_catalog(self):
        # Filter catalog result
        with open(self.search_json_file) as data_file:
            data = json.load(data_file)

        if 'ErrorCode' in data:
            print(data['ErrorMessage'])
            sys.exit(-2)

        # Sort data
        download_dict = {}
        storage_dict = {}
        size_dict = {}
        if len(data["features"]) > 0:
            for i in range(len(data["features"])):
                prod = data["features"][i]["properties"]["productIdentifier"]
                # print(prod, data["features"][i]["properties"]["storage"]["mode"])
                feature_id = data["features"][i]["id"]
                try:
                    storage = data["features"][i]["properties"]["storage"]["mode"]
                    platform = data["features"][i]["properties"]["platform"]
                    resourceSize = int(data["features"][i]["properties"]["resourceSize"])
                    if storage == "unknown":
                        print('found a product with "unknown" status : %s' % prod)
                        print("product %s cannot be downloaded" % prod)
                        print('please send and email with product name to peps admin team : exppeps@cnes.fr')
                    else:
                        # recup du numero d'orbite
                        orbitN = data["features"][i]["properties"]["orbitNumber"]
                        if platform == 'S1A':
                            # calcul de l'orbite relative pour Sentinel 1A
                            relativeOrbit = ((orbitN - 73) % 175) + 1
                        elif platform == 'S1B':
                            # calcul de l'orbite relative pour Sentinel 1B
                            relativeOrbit = ((orbitN - 27) % 175) + 1

                        if self.orbit is not None:
                            if platform.startswith('S2'):
                                if prod.find("_R%03d" % self.orbit) > 0:
                                    download_dict[prod] = feature_id
                                    storage_dict[prod] = storage
                                    size_dict[prod] = resourceSize

                            elif platform.startswith('S1'):
                                if relativeOrbit == self.orbit:
                                    download_dict[prod] = feature_id
                                    storage_dict[prod] = storage
                                    size_dict[prod] = resourceSize
                        else:
                            download_dict[prod] = feature_id
                            storage_dict[prod] = storage
                            size_dict[prod] = resourceSize

                except:
                    pass

            # cloud cover criterium:
            if self.collection[0:2] == 'S2':
                for i in range(len(data["features"])):
                    prod = data["features"][i]["properties"]["productIdentifier"]
                    if data["features"][i]["properties"]["cloudCover"] > self.clouds:
                        del download_dict[prod], storage_dict[prod], size_dict[prod]

            # selecion of specific satellite
            if self.sat != None:
                for i in range(len(data["features"])):
                    prod = data["features"][i]["properties"]["productIdentifier"]
                    if data["features"][i]["properties"]["platform"] != self.sat:
                        try:
                            del download_dict[prod], storage_dict[prod], size_dict[prod]
                        except KeyError:
                            pass

            for prod in download_dict.keys():
                print(prod, storage_dict[prod])
        else:
            print(">>> no product corresponds to selection criteria")
            sys.exit(-1)
        #    print(download_dict.keys())

        return (prod, download_dict, storage_dict, size_dict)

    # sentinel data download function
    def downloader(self):
        if self.search_json_file is None or self.search_json_file == "":
            self.search_json_file = 'search.json'

        if self.sat != None:
            print(self.sat, self.collection[0:2])
            if not self.sat.startswith(self.collection[0:2]):
                print("input parameters collection and satellite are incompatible")
                sys.exit(-1)

        if self.tile is None:
            if self.location is None:
                if self.lat is None or self.lon is None:
                    if (self.latmin is None) or (self.lonmin is None) or (self.latmax is None) or (
                            self.lonmax is None):
                        print("provide at least a point or rectangle or tile number")
                        sys.exit(-1)
                    else:
                        geom = 'rectangle'
                else:
                    if (self.latmin is None) and (self.lonmin is None) and (self.latmax is None) and (
                            self.lonmax is None):
                        geom = 'point'
                    else:
                        print("please choose between point and rectangle, but not both")
                        sys.exit(-1)
            else:
                if (self.latmin is None) and (self.lonmin is None) and (self.latmax is None) and (
                        self.lonmax is None) and (self.lat is None) or (self.lon is None):
                    geom = 'location'
                else:
                    print("please choose location and coordinates, but not both")
                    sys.exit(-1)

        # geometric parameters of catalog request

        if self.tile is not None:
            if self.tile.startswith('T') and len(self.tile) == 6:
                tileid = self.tile[1:6]
            elif len(self.tile) == 5:
                tileid = self.tile[0:5]
            else:
                print("tile name is ill-formated : 31TCJ or T31TCJ are allowed")
                sys.exit(-4)
            query_geom = "tileid=%s" % (tileid)
        elif geom == 'point':
            query_geom = 'lat=%f\&lon=%f' % (self.lat, self.lon)
        elif geom == 'rectangle':
            query_geom = 'box={lonmin},{latmin},{lonmax},{latmax}'.format(
                latmin=self.latmin, latmax=self.latmax, lonmin=self.lonmin, lonmax=self.lonmax)
        elif geom == 'location':
            query_geom = "q=%s" % self.location

        # date parameters of catalog request
        if self.start_date is not None:
            start_date = self.start_date
            if self.end_date is not None:
                end_date = self.end_date
            else:
                end_date = date.today().isoformat()

        # special case for Sentinel-2

        if self.collection == 'S2':
            if self.start_date >= '2016-12-05':
                print("**** products after '2016-12-05' are stored in Tiled products collection")
                print("**** please use option -c S2ST")
            elif self.end_date >= '2016-12-05':
                print("**** products after '2016-12-05' are stored in Tiled products collection")
                print("**** please use option -c S2ST to get the products after that date")
                print("**** products before that date will be downloaded")

        if self.collection == 'S2ST':
            if self.end_date < '2016-12-05':
                print("**** products before '2016-12-05' are stored in non-tiled products collection")
                print("**** please use option -c S2")
            elif self.start_date < '2016-12-05':
                print("**** products before '2016-12-05' are stored in non-tiled products collection")
                print("**** please use option -c S2 to get the products before that date")
                print("**** products after that date will be downloaded")

        # ====================
        # read authentification file
        # ====================
        try:
            f = open(self.auth)
            (email, passwd) = f.readline().split(' ')
            if passwd.endswith('\n'):
                passwd = passwd[:-1]
            f.close()
        except:
            print("error with password file")
            sys.exit(-2)

        if os.path.exists(self.search_json_file):
            os.remove(self.search_json_file)

        # ====================
        # search in catalog
        # ====================
        if (self.product_type == "") and (self.sensor_mode == ""):
            search_catalog = 'curl -k -o %s https://peps.cnes.fr/resto/api/collections/%s/search.json?%s\&startDate=%s\&completionDate=%s\&maxRecords=500' % (
                self.search_json_file, self.collection, query_geom, start_date, end_date)
        else:
            search_catalog = 'curl -k -o %s https://peps.cnes.fr/resto/api/collections/%s/search.json?%s\&startDate=%s\&completionDate=%s\&maxRecords=500\&productType=%s\&sensorMode=%s' % (
                self.search_json_file, self.collection, query_geom, start_date, end_date, self.product_type,
                self.sensor_mode)

        if self.windows:
            search_catalog = search_catalog.replace('\&', '^&')

        print(search_catalog)
        os.system(search_catalog)
        time.sleep(5)

        prod, download_dict, storage_dict, size_dict = self.parse_catalog()

        # ====================
        # Download
        # ====================

        if len(download_dict) == 0:
            print("No product matches the criteria")
        else:
            # first try for the products on tape
            if self.write_dir == None:
                self.write_dir = os.getcwd()

            for prod in list(download_dict.keys()):
                file_exists = os.path.exists(("%s/%s.SAFE") % (self.write_dir, prod)
                                             ) or os.path.exists(("%s/%s.zip") % (self.write_dir, prod))
                if (not (self.no_download) and not (file_exists)):
                    if storage_dict[prod] == "tape":
                        tmticks = time.time()
                        tmpfile = ("%s/tmp_%s.tmp") % (self.write_dir, tmticks)
                        print("\nStage tape product: %s" % prod)
                        get_product = 'curl -o %s -k -u "%s:%s" https://peps.cnes.fr/resto/collections/%s/%s/download/?issuerId=peps &>/dev/null' % (
                            tmpfile, email, passwd, self.collection, download_dict[prod])
                        os.system(get_product)
                        if os.path.exists(tmpfile):
                            os.remove(tmpfile)

            NbProdsToDownload = len(list(download_dict.keys()))
            print("##########################")
            print("%d  products to download" % NbProdsToDownload)
            print("##########################")
            while (NbProdsToDownload > 0):
                # redo catalog search to update disk/tape status
                if (self.product_type == "") and (self.sensor_mode == ""):
                    search_catalog = 'curl -k -o %s https://peps.cnes.fr/resto/api/collections/%s/search.json?%s\&startDate=%s\&completionDate=%s\&maxRecords=500' % (
                        self.search_json_file, self.collection, query_geom, start_date, end_date)
                else:
                    search_catalog = 'curl -k -o %s https://peps.cnes.fr/resto/api/collections/%s/search.json?%s\&startDate=%s\&completionDate=%s\&maxRecords=500\&productType=%s\&sensorMode=%s' % (
                        self.search_json_file, self.collection, query_geom, start_date, end_date,
                        self.product_type, self.sensor_mode)

                if self.windows:
                    search_catalog = search_catalog.replace('\&', '^&')

                os.system(search_catalog)
                time.sleep(2)

                prod, download_dict, storage_dict, size_dict = self.parse_catalog()

                NbProdsToDownload = 0
                # download all products on disk
                for prod in list(download_dict.keys()):
                    file_exists = os.path.exists(("%s/%s.SAFE") % (self.write_dir, prod)
                                                 ) or os.path.exists(("%s/%s.zip") % (self.write_dir, prod))
                    if (not (self.no_download) and not (file_exists)):
                        if storage_dict[prod] == "disk":
                            tmticks = time.time()
                            tmpfile = ("%s/tmp_%s.tmp") % (self.write_dir, tmticks)
                            print("\nDownload of product : %s" % prod)
                            get_product = 'curl -o %s -k -u "%s:%s" https://peps.cnes.fr/resto/collections/%s/%s/download/?issuerId=peps' % (
                                tmpfile, email, passwd, self.collection, download_dict[prod])
                            print(get_product)
                            os.system(get_product)
                            # check binary product, rename tmp file
                            if not os.path.exists(("%s/tmp_%s.tmp") % (self.write_dir, tmticks)):
                                NbProdsToDownload += 1
                            else:
                                self.check_rename(tmpfile, size_dict[prod], prod)

                    elif file_exists:
                        print("%s already exists" % prod)

                # download all products on tape
                for prod in list(download_dict.keys()):
                    file_exists = os.path.exists(("%s/%s.SAFE") % (self.write_dir, prod)
                                                 ) or os.path.exists(("%s/%s.zip") % (self.write_dir, prod))
                    if (not (self.no_download) and not (file_exists)):
                        if storage_dict[prod] == "tape" or storage_dict[prod] == "staging":
                            NbProdsToDownload += 1

                if NbProdsToDownload > 0:
                    print("##############################################################################")
                    print("%d remaining products are on tape, lets's wait 1 minute before trying again" %
                          NbProdsToDownload)
                    print("##############################################################################")
                    time.sleep(60)
